package scu.maqiang.numeric;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Optional;
import java.util.stream.IntStream;

import static java.lang.Math.*;

public class SIMatrix {

	private int m = 0;
	private int n = 0;
	private int nnz = 0;
	private boolean sorted = false;
	public ArrayList<SIEntry>[] mat;

	/**
	 * ϡ������캯�������þ��������������������ʼΪ�վ���
	 * @param m ��������
	 * @param n ��������
	 */
	@SuppressWarnings("unchecked")
	public SIMatrix(int m, int n) {
		this.m = m;
		this.n = n;
		this.nnz = 0;
		mat = new ArrayList[m];
		for (int i = 0; i < m; i++) {
			mat[i] = new ArrayList<SIEntry>();
		}
	}

	/**
	 * ���ξ����캯�������÷��ξ������
	 * @param n ���ξ������
	 */
	public SIMatrix(int n) {
		this(n, n);
	}

	/**
	 * ������ָ�깹��ϡ�����
	 * @param i ������ָ��
	 * @param j ������ָ��
	 * @param value ����Ԫ��ֵ
	 * @param row ��������
	 * @param col ��������
	 */
	public SIMatrix(int[] i, int[] j, int[] value, int row, int col) {
		this(row, col);
		for (int ii = 0, nn = i.length; ii < nn; ii++) {
			if (abs(value[ii]) > Constants.Er) {
				mat[i[ii]].add(new SIEntry(j[ii], value[ii]));
				this.nnz++;
			}
		}
	}

	/**
	 * ������ָ�깹�췽��ϡ�����
	 * @param i ������ָ��
	 * @param j ������ָ��
	 * @param value ����Ԫ��ֵ
	 * @param size �������
	 */
	public SIMatrix(int[] i, int[] j, int[] value, int size) {
		this(i, j, value, size, size);
	}

	/**
	 * ����֪�������¾���
	 * @param sm ��֪����
	 */
	@SuppressWarnings("unchecked")
	public SIMatrix(SIMatrix sm) {
		this(sm.m, sm.n);
		this.m = sm.m;
		this.n = sm.n;
		mat = new ArrayList[m];
		for (int i = 0; i < this.m; i++) {
			mat[i] = new ArrayList<SIEntry>();
			for (SIEntry ele : sm.mat[i]) {
				mat[i].add(new SIEntry(ele));
				this.nnz++;
			}
		}
	}

	/**
	 * �õ�����Ԫ�ص�ֵ
	 * @param i Ԫ����ָ��
	 * @param j Ԫ����ָ��
	 * @return ������ڣ��򷵻�Ԫ�ص�ֵ�����򷵻�0
	 */
	public int getElement(int i, int j) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		Optional<SIEntry> element = mat[i].stream().filter(ele -> ele.col == j).findFirst();
		return element.isPresent() ? element.get().val : 0;
	}

	/**
	 * ���þ���Ԫ�ص�ֵ
	 * @param i Ԫ����ָ��
	 * @param j Ԫ����ָ��
	 * @param value Ԫ��ֵ
	 * @return ��ǰϡ�����
	 */
	public SIMatrix setElement(int i, int j, int value) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		int size = mat[i].size();
		if (size == 0 && abs(value) > Constants.Er) {
			mat[i].add(new SIEntry(j, value));
			this.nnz++;
			return this;
		}

		for (int idx = 0; idx < size; idx++) {
			SIEntry se = mat[i].get(idx);
			if (se.col == j) {
				if (abs(value) > Constants.Er)
					se.val = value;
				else {
					mat[i].remove(idx);
					--(this.nnz);
				}
				return this;
			}
		}

		mat[i].add(new SIEntry(j, value));
		++(this.nnz);
		return this;
	}

	/**
	 * �����κμ�����þ���Ԫ�ص�ֵ
	 * @param i Ԫ����ָ��
	 * @param j Ԫ����ָ��
	 * @param value Ԫ��ֵ
	 * @return ��ǰϡ�����
	 */
	public SIMatrix setElementFast(int i, int j, int value) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		mat[i].add(new SIEntry(j, value));
		++(this.nnz);
		return this;
	}

	/**
	 * ��ϡ�������Ӧ����ָ���ϼ�����ӦԪ��ֵ, ���ս�������������:
	 * 1. ���Ԫ��ֵΪ0�������κβ���
	 * 2. ���Ԫ��ֵ��0�����Ҿ���������ָ���д���Ԫ�أ���ԭ��Ԫ��ֵ��֮���,
	 *    ���Ӻ����Ϊ0, ��ɾ����Ԫ�أ������0�����޸ĸ�Ԫ��ֵ
	 * 3. ���Ԫ��ֵ��0���������в����ڸ�Ԫ�أ���������������Ӹ�Ԫ��ֵ
	 * @param i Ԫ����ָ��
	 * @param j Ԫ����ָ��
	 * @param val Ԫ��ֵ
	 *            
	 */
	public void addElement(int i, int j, int val) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		if (abs(val) > Constants.Er) {
			int size = mat[i].size();
			if (size == 0) {
				mat[i].add(new SIEntry(j, val));
				this.nnz++;
				return;
			}
			for (int idx = 0; idx < size; idx++) {
				SIEntry se = mat[i].get(idx);
				if (se.col == j) {
					int temp = se.val + val;
					if (temp == 0) {
						mat[i].remove(idx);
						--(this.nnz);
					}else {
						se.val = temp;
					}
					return;
				}
			}
			mat[i].add(new SIEntry(j, val));
			++(this.nnz);
			return;
		}
	}

	/**
	 * �ھ�����ɾ����i�е�i������Ԫ��ֵ���������ϡ����������з�0Ԫ�ص�����ָ���ǶԳƵ�
	 * @param i ��ָ��
	 */
	public void removeSymmetricIJ(int i) {
		mat[i].forEach(ele -> {
			int col = ele.col;
			if (col != i) {
				removeElement(col, i);
			}
		});
		removeRow(i);
	}

	/**
	 * ȥ��ϡ�������ĳ����ָ���µ�Ԫ��
	 * @param i Ԫ����ָ��
	 * @param j Ԫ����ָ��
	 * @return ���ɾ����Ԫ�ش��ڣ��򷵻�true; ���򷵻�false
	 */
	public boolean removeElement(int i, int j) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		int size = mat[i].size();
		for (int idx = 0; idx < size; idx++) {
			if (mat[i].get(idx).col == j) {
				mat[i].remove(idx);
				--(this.nnz);
				return true;
			}
		}
		return false;
	}

	/**
	 * ɾ��ϡ�����ĳ������Ԫ��
	 * @param i ������ָ��
	 */
	public void removeRow(int i) {
		ParamCheck.checkBounds(i, 0, m);
		this.nnz -= mat[i].size();
		mat[i].clear();
	}

	/**
	 * �õ���������
	 * @return ��������
	 */
	public int getM() {
		return m;
	}

	/**
	 * �õ���������
	 * @return ��������
	 */
	public int getN() {
		return n;
	}

	/**
	 * �õ������0Ԫ�ظ���
	 * @return �����0Ԫ�ظ���
	 */
	public int getNNZ() {
		return nnz;
	}

	public int getNNZL() {
		int sum = 0;
		for(int i = 0; i < m; i++) {
			int ii = i;
			sum += mat[i].stream().filter(ele -> ele.col < ii).count();
		}
		return sum;
	}
	
	public int getNNZU() {
		int sum = 0;
		for(int i = 0; i < m; i++) {
			int ii = i;
			sum += mat[i].stream().filter(ele -> ele.col > ii).count();
		}
		return sum;		
	}
	
	public Pair<Integer, Integer> getNNZLU() {
		int numL = 0;
		int numU = 0;
		for(int i = 0; i < m; i++) {
			int ii = i;
			numL += mat[i].stream().filter(ele -> ele.col < ii).count();
			numU += mat[i].stream().filter(ele -> ele.col < ii).count();
		}
		return new Pair<Integer, Integer>(numL, numU);
	}
	
	
	/**
	 * ʵ��ϡ��������Ԫ��, �����ǰ����ΪA, �������Ϊalpha��B, ���յ÷��ؾ���ֵCΪ
	 *    C �� A + alpha * B.
	 * @param alpha  ʵ��
	 * @param B����Ӿ���
	 * @return�����ؾ���ֵΪC = A + alpha * B
	 */
	public SIMatrix add(int alpha, SIMatrix B) {
		SIMatrix result = new SIMatrix(this);
		int M = result.m;
		boolean flag = false;
		if (abs(alpha) < Constants.Er) {
			return result;
		}

		for (int i = 0; i < M; i++) {
			int mat2Size = B.mat[i].size();
			for (int j = 0; j < mat2Size; j++) {
				// flag��¼�Ƿ��ڵڶ��������д��ڸ�ָ��
				flag = false;
				SIEntry e2 = B.mat[i].get(j);
				int resultSize = result.mat[i].size();
				for (int k = 0; k < resultSize; k++) {
					SIEntry e1 = result.mat[i].get(k);
					// Ѱ���������е���ͬ�����Ƿ�����ͬ����ָ��
					if (e1.col == e2.col) {
						int temp2 = e1.val + alpha * e2.val;
						// �������ͬ��ָ�꣬����֮������ֲ�Ϊ0,���޸���ָ����ֵ
						if (temp2 == 0) {
							e1.val = temp2;
						} else {
							// ������Ϊ0,������Ԫ��
							result.mat[i].remove(k);
							result.nnz--;
						}
						flag = true;
						break;
					}
				}
				// ���û�и�ָ�꣬��ֱ�Ӽ����Ԫ��
				if (!flag) {
					result.mat[i].add(new SIEntry(e2));
					result.nnz++;
				}
			}
		}
		return result;
	}

	/**
	 * 
	 * @param d
	 * @return
	 */
	public SIMatrix mul(int d) {
		SIMatrix result = new SIMatrix(m, n);
		mmul(d, result);
		return result;
	}

	
	public void mul(int d, SIMatrix result) {
		ParamCheck.checkEqual(result.m, m);
		ParamCheck.checkEqual(result.n, n);
		mmul(d, result);
	}
	
	private void mmul(int d, SIMatrix result) {
		IntStream.range(0, m).forEach(i -> mat[i].forEach(ele -> {
			result.mat[i].add(new SIEntry(ele.col, ele.val * d));
			result.nnz++;
		}));
		/*
		for (int i = 0; i < m; i++) {
			for (SIEntry ele : mat[i]) {
				result.mat[i].add(new SIEntry(ele.col, ele.val * d));
				result.nnz++;
			}
		}*/		
	}

	public double[] mul(double[] u) {
		ParamCheck.checkEqual(n, u.length);
		double[] v = new double[m];
		mmul(u, v);
		return v;
	}

	public void mul(double[] u, double[] v) {
		ParamCheck.checkEqual(n, u.length);
		ParamCheck.checkEqual(m, v.length);
		mmul(u, v);
	}
	
	private void mmul(double[] u, double[] v) {
		//Arrays.parallelSetAll(v, i -> mat[i].parallelStream().mapToDouble(ele -> ele.val * u[ele.col]).sum());
		for(int i = 0; i < m; i++) {
			v[i] = 0.0;
			for(SIEntry ele: mat[i]) {
				v[i] += ele.val * u[ele.col];
			}
		}
	}
	
	public double[] mulT(double[] u) {
		ParamCheck.checkEqual(m, u.length);
		double[] v = new double[n];
		mmulT(u, v);
		return v;
	}
	
	
	public void mulT(double[] u, double[] v) {
		ParamCheck.checkEqual(m, u.length);
		ParamCheck.checkEqual(n, v.length);
		mmulT(u, v);		
	}
	/**
	 * 
	 * @param u
	 * @param v
	 */
	private void mmulT(double[] u, double[] v) {
		Arrays.fill(v, 0.0);
		for(int i = 0; i < m; i++) {
			for(SIEntry ele: mat[i]) {
				v[ele.col] += ele.val * u[i];
			}
		}
	}

	/**
	 * 
	 * @param rowDof
	 * @param colDof
	 * @param ele
	 */
	public void assemble(int[] rowDof, int[] colDof, int[][] ele) {
		int m = ele.length;
		int n = ele[0].length;
		ParamCheck.checkEqual(m, rowDof.length);
		ParamCheck.checkEqual(n, colDof.length);
		for (int i = 0; i < m; i++) {
			for (int j = 0; j < n; j++) {
				addElement(rowDof[i], colDof[j], ele[i][j]);
			}
		}
		
		
	}

	/**
	 * 
	 * @param dof
	 * @param ele
	 */
	public void assemble(int[] dof, int[][] ele) {
		assemble(dof, dof, ele);
	}

	public void assemble(int row, int[] colDof, int[] ele) {
		int m = ele.length;
		ParamCheck.checkEqual(m, colDof.length);
		for (int i = 0; i < m; i++) {
			addElement(row, colDof[i], ele[i]);
		}
	}

	/**
	 * 
	 * @param rowDof
	 * @param col
	 * @param ele
	 */
	public void assemble(int[] rowDof, int col, int[] ele) {
		int m = ele.length;
		ParamCheck.checkEqual(m, rowDof.length);
		for (int i = 0; i < m; i++) {
			addElement(rowDof[i], col, ele[i]);
		}
	}

	/**
	 * 
	 */
	public void sort() {
		//for (int i = 0; i < m; i++) {
		//	Collections.sort(mat[i], (a, b) -> a.col > b.col ? 1 : -1);
		//}
		IntStream.range(0, m).parallel().forEach(i -> Collections.sort(mat[i], (a, b) -> a.col > b.col ? 1 : -1));
		sorted = true;
	}

	/**
	 * 
	 * @return
	 */
	public boolean isSorted() {
		return sorted;
	}
	
	/**
	 * 
	 * @return
	 */
	public boolean makeEntrySymmetric() {
		int oldNNZ = nnz;
		for (int i = 0; i < m; i++) {
			int len = mat[i].size();
			for (int j = 0; j < len; j++) {
				int col = mat[i].get(j).col;
				if (abs(getElement(col, i)) < Constants.Er) {
					setElementFast(col, i, 0);
				}
			}
		}
		return oldNNZ != nnz;
	}

	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		sb.append(m + "\t" + n + "\t" + nnz + "\t" + isSymmetric() + "\n");
		for (int i = 0; i < m; i++) {
			for (SIEntry ele : mat[i]) {
				sb.append((i + 1) + "\t" + (ele.col + 1) + "\t" + ele.val + "\n");
			}

		}
		return sb.toString();
	}

	public double[] diag() {
		double[] d = new double[m];
		for (int i = 0; i < m; i++) {
			for (SIEntry ele : mat[i]) {
				if (ele.col == i) {
					d[i] = ele.val;
				}
			}
		}
		return d;
	}

	public SIMatrix transpose() {
		SIMatrix result = new SIMatrix(n, m);
		for (int i = 0; i < m; i++) {
			for (SIEntry ele : mat[i]) {
				result.mat[ele.col].add(new SIEntry(i, ele.val));
			}
		}
		return null;
	}

	public SIMatrix triU() {
		SIMatrix upper = new SIMatrix(m, n);
		for (int i = 0; i < m; i++) {
			int ii = i;
			mat[i].stream().filter(ele -> ele.col > ii).forEach(ele -> upper.setElementFast(ii, ele.col, ele.val));
		}
		return upper;
	}
	
	public SIMatrix triL() {
		SIMatrix lower = new SIMatrix(m, n);
		for (int i = 0; i < m; i++) {
			int ii = i;
			mat[i].stream().filter(ele -> ele.col < ii).forEach(ele -> lower.setElementFast(ii, ele.col, ele.val));
		}
		return lower;
	}
	/**
	 * ʹ����Ԫ�صõ������Ԫ�ص�ֵ
	 * @param ii��������ָ��
	 * @param jj��������ָ��
	 * @param values������Ԫ��ֵ
	 */
	public void getTriValues(int[] ii, int[] jj, int[] values) {
		int in = 0;
		for (int i = 0; i < m; i++) {
			for (SIEntry ele : mat[i]) {
				ii[in] = i;
				jj[in] = ele.col;
				values[in] = ele.val;
				in++;
			}
		}
	}
	
	/**
	 * ����õ�ϡ�����Ԫ�����ֵ
	 * @return��ϡ�����Ԫ�����ֵ
	 */
	public double maxElement() {
		double[] maxval = new double[m];
		for(int i = 0; i < m; i++) {
			maxval[i] = mat[i].stream().mapToDouble(ele -> ele.val).max().orElse(0.0);
		}
		return Arrays.stream(maxval).max().orElse(0.0);
	}
	
	/**
	 * ����õ�ϡ�����Ԫ�����ֵ�����������ֵ��������ֵ
	 * @param IJ�����ֵ��������ֵ
	 * @return��ϡ�����Ԫ�����ֵ
	 */
	public double maxElement(int[] IJ) {
		double[] maxval = new double[m];
		int[] maxJ = new int[m];
		for (int i = 0; i < m; i++) {
			if (mat[i] != null) {
				int tempJ = mat[i].get(0).col;
				double tempVal = mat[i].get(0).val;
				for(int j = 1, size = mat[i].size(); j < size; j++) {
					SIEntry ele = mat[i].get(j);
					if (ele.val > tempVal) {
						tempJ = ele.col;
						tempVal = ele.val;
					}
				}
				maxJ[i] = tempJ;
				maxval[i] = tempVal;
			}
		}
		double result = maxval[0];
		IJ[0] = 0;
		IJ[1] = maxJ[0];
		for (int i = 1; i < m; i++) {
			if (result < maxval[i]) {
				result = maxval[i];
				IJ[0] = i;
				IJ[1] = maxJ[i];
			}
		}
		return result;
	}
	
	public boolean isSymmetric() {
		for(int i = 0; i < m; i++) {
			for (SIEntry ele : mat[i]) {
				if(abs(ele.val - getElement(ele.col, i)) > Constants.Er) {
					return false;
				}
			}
		}
		return true;
	}
}
